package developer.zyk.ability.chat.ability.rag.embedding;

import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import developer.zyk.ability.config.AIProperties;
import developer.zyk.global.config.PlatFormProperties;
import jakarta.annotation.Resource;
import lombok.Data;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author zhangyongkang
 * @time 2025/5/15 10:30
 * @description
 */
@Component("aliEmbeddingModel")
public class AliEmbeddingModel extends AbstractEmbeddingModel {
    @Resource
    private PlatFormProperties platFormProperties;

    @Override
    public EmbeddingResponse call(EmbeddingRequest request) {
        return new EmbeddingResponse(invoke(request.getInstructions()));
    }

    @Override
    public float[] embed(Document document) {
        return invoke(List.of(document.getText())).get(0).getOutput();
    }

    @Override
    public int dimensions() {
        return 1;
    }


    private List<Embedding> invoke(List<String> strings) {
        PlatFormProperties.AliBaiLian aliBaiLian = platFormProperties.getAliBaiLian();
        JSONObject jsonObject = new JSONObject();
        jsonObject.set("model", aliBaiLian.getEmbeddingModel());
        jsonObject.set("input", strings);
        jsonObject.set("encoding_format", "float");
        HttpResponse execute = HttpUtil.createPost(aliBaiLian.getEmbeddingUrl())
                .headerMap(headerMap(aliBaiLian.getApiKey()), true)
                .body(jsonObject.toStringPretty())
                .execute();

        return JSONUtil.toBean(execute.body(), AliEmbeddingResult.class)
                .getData()
                .stream()
                .map(a -> {
                    Embedding embedding = new Embedding(a.getEmbedding(), a.getIndex());
                    return embedding;
                })
                .toList();
    }

    private Map<String, String> headerMap(String key) {
        Map<String, String> header = new HashMap<>();
        header.put("Authorization", " Bearer " + key);
        header.put("Content-Type", "application/json");
        return header;
    }

    @Data
    public static class AliEmbeddingResult {
        private List<Result> data;
        private String object;
        private String model;
        private String id;

        @Data
        public static class Result {
            private float[] embedding;
            private Integer index;
            String object;
        }
    }


}
